#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "PublicDNS.h"
#include "PublicIPTool.h"
#include "PublicSocket.h"

#if	1
#define DEGUG_PRINTF(format, arg...) \
	printf(format, ##arg);
#else
#define DEGUG_PRINTF(format, arg...)  
#endif

//__asm__(".symver __isoc99_sscanf,sscanf@GLIBC_2.2.5");
struct DNS_HEADER{
	unsigned	short id;		    // identification number
	
	unsigned	char rd     :1;		// recursion desired
	unsigned	char tc     :1;		// truncated message
	unsigned	char aa     :1;		// authoritive answer
	unsigned	char opcode :4;	    // purpose of message
	unsigned	char qr     :1;		// query/response flag
	
	unsigned	char rcode  :4;	    // response code
	unsigned	char z      :3;		// its z! reserved
	unsigned	char ra     :1;		// recursion available
	
	unsigned    short q_count;	    // number of question entries
	unsigned	short ans_count;	// number of answer entries
	unsigned	short auth_count;	// number of authority entries
	unsigned	short add_count;	// number of resource entries
};

struct DNS_QUESTION{
	unsigned short qtype;
	unsigned short qclass;
};

#pragma pack(push, 1)
struct  DNS_RDATA{
	unsigned short type;
	unsigned short _class;
	unsigned int   ttl;
	unsigned short data_len;
};
#pragma pack(pop)

struct DNS_RECORD{
	unsigned char		*name;
	struct DNS_RDATA	*resource;
	unsigned char		*rdata;
};

int ListInitiate(SLNode **L){
	if((*L=(SLNode *)malloc(sizeof(SLNode)))==NULL)
		return -1;
	(*L)->next=NULL;
		return 0;
}

int ListInsert(SLNode **L,recordinfo *x){
	SLNode *p,*q;
	p=*L;
	if(p == NULL)return -1;
	while(p->next !=NULL){
		p=p->next;
	}
	if((q=(SLNode *)malloc(sizeof(SLNode)))==NULL)
		return -1;
	memcpy(&q->data,x,sizeof(recordinfo));  	
	q->next=p->next;
	p->next=q;
	return 0;
}

void Destroy(SLNode * *L){
	SLNode *p,*p1;
	p=*L;
	while(p!=NULL){
		p1=p;
		p=p->next;
		free(p1);
	}
	*L=NULL;
}

void changetoDnsNameFormat(unsigned char* dns,unsigned char* host){
	int lock=0, i = 0;

	if(host[strlen((char*)host) -1] != '.')
		strcat((char*)host,".");

	for(i=0;i<(int)strlen((char*)host);i++){
		if(host[i]=='.'){
			*dns++= i-lock;
			for(;lock<i;lock++){
				*dns++ = host[lock];
			}
			lock++;
		}
	}
	*dns++='\0';
}

void recoverFromDnsNameFormat(const unsigned char* dns,unsigned char* host){
	int i , j;
	unsigned int p=0;
	for(i=0;i<(int)strlen((char*)dns);i++){
		p=dns[i];
		for(j=0;j<(int)p;j++){
			host[i]=dns[i+1];
			i=i+1;
		}
		host[i]='.';
	}
	return;	
}

unsigned char*readDnsName(unsigned char*reader,unsigned char*buffer,int*count){
	unsigned char *name;
	unsigned int p=0,jumped=0,offset;
	int i , j;
	
	*count = 1;
	name   = (unsigned char*)malloc(512);
	
	name[0]='\0';
	//read the names in 3www6google3com format
	while(*reader!=0){
		if(*reader>=192){
			offset = (*reader)*256 + *(reader+1) - 49152; //49152 = 11000000 00000000  ;)
#if 0
			reader = buffer + offset - 1;
			jumped = 1;  //we have jumped to another location so counting wont go up!
#else
			if(offset!=(uint32_t)(reader-buffer)){
				reader = buffer + offset - 1;
				jumped = 1;  //we have jumped to another location so counting wont go up!
			}else{
//新定位位置与当前一致情况直接跳过
				reader+=2;
				jumped =1;
				break;
			}
#endif
		}else{
			name[p++]=*reader;
		}
		reader=reader+1;
		if(jumped==0) *count = *count + 1; //if we havent jumped to another location then we can count up
	}
	
	name[p]='\0';    //string complete
	if(jumped==1) *count = *count + 1;  //number of steps we actually moved forward in the packet
	
	//now convert 3www6google3com0 to www.google.com
	for(i=0;i<(int)strlen((const char*)name);i++){
		p=name[i];
		for(j=0;j<(int)p;j++){
			name[i]=name[i+1];
			i=i+1;
		}
		name[i]='.';
	}
	return name;		
}
int parse_dns_query(CPublicSocket*sock,const int times,const char*domain,std::set<std::string>iplist,SLNode**ppstRecordInfo,int iptype){	
	char szBuffer[1024];
	char domainStr[256]={0};
	int ret = 0, len = 0;

	unsigned char* qname = NULL;
	struct DNS_HEADER* dns = NULL;
	struct DNS_QUESTION* qinfo = NULL;

	char buf[1024]={0};

	int nindex = 0, i = 0, nStop = 0;
	char chHost[128] = {0};

	struct DNS_RECORD record;
	struct DNS_HEADER* header = NULL;
//	struct DNS_QUESTION* question = NULL;
	unsigned char* pQueries = NULL;
	unsigned char* pAnswers = NULL;
//	struct sockaddr_in a;	

	int tryCnt = 0;
	struct  timeval  start;
	struct  timeval  end;
	unsigned long timer = 0;
	int count = 0;
	int bGetDns=0;

	DEGUG_PRINTF("Request domain %s\n", domain);

	memset(szBuffer,0x0,sizeof(szBuffer));
	dns = (struct DNS_HEADER*)&szBuffer;
	strcpy(domainStr, domain);

	srand(time(NULL));
//	id=(unsigned short)htons(random(65535) + 1);
//	id2=id1+1;
	dns->id = hton16(random(256));
	dns->qr = 0;dns->opcode = 0;dns->aa = 1;dns->tc = 0;dns->rd = 1;dns->ra = 0;dns->z = 0;dns->rcode = 0;
	dns->q_count = hton16(1);
	dns->ans_count = dns->auth_count = dns->add_count = 0;
	qname = (unsigned char*)&szBuffer[sizeof(struct DNS_HEADER)];
	changetoDnsNameFormat(qname, (unsigned char*)domainStr);
	
	qinfo  = (struct DNS_QUESTION*)&szBuffer[sizeof(struct DNS_HEADER) + (strlen((const char*)qname) + 1)];
	//type 1ipv4,28ipv6
	if(iptype==6){
		qinfo->qtype = hton16(28);
	}else{
		qinfo->qtype = hton16(1);
	}
	//class 1 in
	qinfo->qclass = hton16(1); 
	len = sizeof(struct DNS_HEADER) + (strlen((const char*)qname) + 1) + sizeof(struct DNS_QUESTION);

	ret = ListInitiate(ppstRecordInfo);

	memset(&start, 0x0, sizeof(struct timeval));
	gettimeofday(&start,NULL);
	std::set<std::string>::iterator it=iplist.begin();
	while(it!=iplist.end()){
		ret=sock->Send(szBuffer,len,it->c_str(),53);	
		DEGUG_PRINTF("SendTo:%s\n",it->c_str());
//		if(ret == -1){
//			DEGUG_PRINTF("send dns query fail, errno=%d\n", errno);
//			return -1;
//		}
		it++;
	}
	char FromIP[40]={0};
recv_dns:
	do{
		memset(FromIP,0,sizeof(FromIP));
		memset(&end, 0x0, sizeof(struct timeval));
		ret=sock->Recv(buf,sizeof(buf),FromIP,&end);		
	}while(0);	
	if(ret >0){
		int buflen=ret;
		header = (struct DNS_HEADER*)buf;
//不是想要的ID,重试5次
		if(header->id!=dns->id){
			if(tryCnt++ < 5)
				goto recv_dns;
			else
				goto end_dns;
		}
		if(header->qr == 0 || hton16(header->ans_count) == 0){
			if(iplist.size()>1)
				goto recv_dns;
			else
				goto end_dns;
		}
		pQueries = (unsigned char*)buf+ sizeof(struct DNS_HEADER);
		if(pQueries == NULL || *pQueries == 0)
			goto recv_dns;
		recoverFromDnsNameFormat(pQueries, (unsigned char*)domainStr);
		nindex = strlen((const char*)pQueries) +1;
//		question =  (struct DNS_QUESTION*)&pQueries[nindex];
		pAnswers = pQueries+nindex+sizeof(struct DNS_QUESTION);
		if(pAnswers == NULL || *pAnswers == 0){
			DEGUG_PRINTF("[%s] can't get result from [%s], error[%d]\n", domainStr, FromIP, header->rcode);
			goto recv_dns;
		}
//		DEGUG_PRINTF("header->ans_count is %d\n", ntohs(header->ans_count));
		recordinfo *pstTemp = NULL;
		pstTemp = (recordinfo *)malloc(sizeof(recordinfo));
		if(NULL == pstTemp)	goto recv_dns;
		memset(pstTemp, 0x0, sizeof(recordinfo));

		DEGUG_PRINTF("timer begin= %ld s %ld us\n",start.tv_sec,start.tv_usec);
		DEGUG_PRINTF("timer end= %ld s %ld us\n",end.tv_sec,end.tv_usec);	
		timer = 1000000*(end.tv_sec-start.tv_sec)+(end.tv_usec-start.tv_usec);
		DEGUG_PRINTF("timer = %ld us\n",timer);		
		
		for(i = 0; i < hton16(header->ans_count); i ++){
			record.name = readDnsName(pAnswers, (unsigned char*)buf,&nStop);
			//first answer may be CNAME record
			//不要区分大小写
			if(i == 0 && strcasecmp((const char*)record.name, domainStr) != 0){
				DEGUG_PRINTF("Get wrong dns response [%s] with same id, continue recving...\n", record.name);
				free(record.name);
				if(tryCnt++ < 5){
					count = 0;
					free(pstTemp);
					goto recv_dns;
				}else{
					free(pstTemp);
					goto end_dns;
				}
			}
			pAnswers = pAnswers + nStop;
			record.resource = (struct DNS_RDATA*)(pAnswers);
			pAnswers = pAnswers + sizeof(struct DNS_RDATA);
//			DEGUG_PRINTF("Get type %d\n", ntohs(record.resource->type));
//			if its an ipv4 address
			if(hton16(record.resource->type) == 1) {
				//a.sin_addr.s_addr = *(unsigned  int*)pAnswers;
				//memset(chHost,0,sizeof(chHost));
				//strncpy(chHost,inet_ntoa(a.sin_addr),strlen(inet_ntoa(a.sin_addr)));				
				GetIPByHostID(chHost, *(uint32_t*)pAnswers);
				pAnswers = pAnswers + hton16(record.resource->data_len);
//				DEGUG_PRINTF("Get A record: %s for %s, TTL %u\n", chHost, record.name, ntohl(record.resource->ttl));
				pstTemp->ttl = hton32(record.resource->ttl);
				pstTemp->queryCostTime = timer;
				pstTemp->queryCnt = times;
				strcpy(pstTemp->host,chHost);
				strcpy(pstTemp->server,FromIP);
				DEGUG_PRINTF("Get A record: host %s, count %d, TTL %u\n",pstTemp->host,count+1, pstTemp->ttl);
				ret = ListInsert(ppstRecordInfo,pstTemp);
				bGetDns=1;
				count ++;
			}else
//ipv6				
			if(hton16(record.resource->type) == 28) {
				GetIPByHostIDV6(chHost,(uint8_t*)pAnswers);
				pAnswers = pAnswers + hton16(record.resource->data_len);
//				DEGUG_PRINTF("Get A record: %s for %s, TTL %u\n", chHost, record.name, ntohl(record.resource->ttl));
				pstTemp->ttl = hton32(record.resource->ttl);
				pstTemp->queryCostTime = timer;
				pstTemp->queryCnt = times;
				strcpy(pstTemp->host,chHost);
				strcpy(pstTemp->server,FromIP);
				DEGUG_PRINTF("Get A record: host %s, count %d, TTL %u\n",pstTemp->host,count+1, pstTemp->ttl);
				ret = ListInsert(ppstRecordInfo,pstTemp);
				bGetDns=1;
				count ++;
			}else{
				pAnswers = pAnswers + hton16(record.resource->data_len);
			}
			free(record.name);
//是否越界判断
			if((char*)pAnswers>(buf+buflen)){
				break;
			}
		}
		free(pstTemp);
//未取到并且数量大于1，重新获取DNS
		if((bGetDns==0)&&(iplist.size()>1)){
			 goto recv_dns;
		}
	}else{
//没有数据返回直接结束		
		goto end_dns;
	}
end_dns:
	if(bGetDns)return 0;
	return -1;
}
int getNameServerCfg(SETIPLIST&iplist){
	FILE* fs = NULL;
	char line[256] = {0};
	char *p = NULL;
	int success = 0;
	fs = fopen("/etc/resolv.conf", "r");
	if(fs == NULL)	return -1;

	while(fgets(line, sizeof(line), fs)){
		if(strncmp(line, "nameserver", strlen("nameserver")) == 0){
			//get nameserver data
			while(line[strlen(line)-1] == 0xa || line[strlen(line)-1] == 0xd){
		        line[strlen(line)-1] = '\0';
			}
			p = (char*)line + strlen("nameserver");
			while(*p == ' ' || *p == '\t'){
				p++;	
			}
//不是IPV4跳过			
			if(GetIPType(p)!=0){		
				iplist.insert(p);
				success = 1;
			}
		}
		memset(line,0,sizeof(line));
	}
	fclose(fs);
	if(success == 0)	return -1;
	return 0;
}
//从/etc/hosts获取IP
int gethostFromHosts(const char*domain,SLNode**ppstRecordInfo,int iptype){
	 FILE* fs = NULL;
	 char line[1024] = {0};
	 struct  timeval  start={0};
	 gettimeofday(&start,NULL);
	 fs = fopen("/etc/hosts", "r");
	 if(fs==NULL)return -1;
	 while(fgets(line, sizeof(line), fs)){
		char IP[64]={0};
		char Host[128]={0};
//非0～9开头都过滤掉
		if(line[0]<'0'||line[0]>'9'){
			//DEGUG_PRINTF("Pass Line:%s\n",line);
		}else{
			if(sscanf(line," %s %s",IP,Host)==2){
				if(strcmp(Host,domain)==0){
//判段跟期望的IP类型是否一致					
					if(GetIPType(IP)==iptype){				
						struct  timeval  end={0};
						recordinfo info={0};
						gettimeofday(&end,NULL);
						printf("Find IP:%s Host:%s\n",IP,Host);
						info.queryCnt=1;
						info.queryCostTime=1000000*(end.tv_sec-start.tv_sec)+(end.tv_usec-start.tv_usec);
						info.ttl=1;
						info.hosts=1;
						strcpy(info.server,"localhost");
						strcpy(info.host,IP);
						ListInitiate(ppstRecordInfo);
						ListInsert(ppstRecordInfo,&info);
						break;
					}
				}
			}else{
				DEGUG_PRINTF("Pass Line2:%s\n",line);
		    }
		}
		memset(line,0,sizeof(line));
	 }
	 fclose(fs);
	 if(ppstRecordInfo!=NULL)return 0;
	 return -1;
}
SLNode* dlgethostbyname(const char*domain,SETIPLIST dnslist,int timeout,int trytimes,int iptype){
	int  ret = 0;
	int  index = 0;
	SETIPLIST iplist;
	SLNode *pRecordList = NULL;
	DEGUG_PRINTF("enter dlgethostbyname %zd\n",dnslist.size());
	if(NULL == domain || *domain == '\0'){
		printf("invalid domain\n");
		return NULL;
	}
//传入DNS为空	
	if(dnslist.empty()){
//尝试从/etc/hosts中获取IP
		ret = gethostFromHosts(domain,&pRecordList,iptype);
		if(ret==0&&pRecordList!=NULL){
			return pRecordList;
		}
		ret = getNameServerCfg(iplist);
		if(ret == -1){
			DEGUG_PRINTF("getNameServerCfg fail\n");
			return NULL;
		}
	}else{
		SETIPLIST::iterator itdns=dnslist.begin();
		while(itdns!=dnslist.end()){
			if((*itdns)!=""){
				iplist.insert(itdns->c_str());
			}
			itdns++;
		}
	}

	if(iplist.empty()){
		DEGUG_PRINTF("获取DNSIP失败\n");
		return NULL;
	}
	CPublicSocket sock;
	ret=sock.CreateSocket(SOCKET_TYPE_UDP,timeout,NULL);	
	if(ret == -1)return NULL;
	for(index=0;index<trytimes;index++){		
		ret = parse_dns_query(&sock,index+1,domain,iplist,&pRecordList,iptype);		
		if((ret == 0) && (NULL != pRecordList)){
			break;
		}else
		if(pRecordList){
			Destroy(&pRecordList);
		}
	}
	return pRecordList;
}
#ifndef NOUSE_PUBLIC_LOG
#include"PublicLog.h"
#endif
void PrintDNSInfo(SLNode*result,const char*domain){
	SLNode *pTmp = result;
	while(pTmp->next !=NULL){
		pTmp=pTmp->next;
		recordinfo* pRrdata=NULL;
		pRrdata= &pTmp->data;
		INFO("[%s]dns[%s] host=%s,query count=%u,query time=%ld us,localhost:%hhd",
		domain, pRrdata->server, pRrdata->host,pRrdata->queryCnt, pRrdata->queryCostTime,pRrdata->hosts);
	}
}
#if 0
//usage: dnsdig www.qq.com or dnsdig @1.1.1.1 www.qq.com 
int main(int argc, char** argv){
	int ret = 0;
	unsigned long ulTimer = 0;
	SLNode *pRecordList = NULL;
	SLNode *pTmp = NULL;
	char domain[256] = {0}, server[40] = {0};

	recordinfo* pRrdata = NULL;


	if((argc != 2) && (argc != 3)){
		printf("Usage: dnsdig [domain]\nfor example: dnsdig www.qq.com or dnsdig @1.1.1.1 www.qq.com\n");
		return 0;
	}
	if(argc == 3){
		if(argv[1][0] == '@'){
			strcpy(server, argv[1]+1);
			strcpy(domain, argv[2]);
		}else if(argv[2][0] == '@'){
			strcpy(server, argv[2]+1);
			strcpy(domain, argv[1]);
		}
		pRecordList = dlgethostbyname(domain, server);
	}else if(argc == 2){
		strcpy(domain, argv[1]);
		pRecordList = dlgethostbyname(domain, NULL);
	}


	if(NULL == pRecordList)
	{
		printf("dlgethostbyname failed!\n");
		return -1;
	}
	else
	{
		pTmp = pRecordList;
		while(pTmp->next !=NULL)
		{
			pTmp=pTmp->next;
			pRrdata = &pTmp->data;
			printf("[%s] via [%s]: host = %s, ttl = %u, query count = %u,  query cost time=%ums\n", 
				domain, pRrdata->server, pRrdata->host, pRrdata->ttl, pRrdata->queryCnt, pRrdata->queryCostTime/1000);
		}
		Destroy(&pRecordList);
	}
	return 0;
}
#endif

